package com.sprint.questai.module.kb;

/**
 * @Author: mayixiang
 * @Date: 2024/6/14 10:08
 * @Description: 知识库操作的服务
 */

import cn.hutool.core.util.RandomUtil;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.IndexResponse;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import com.sprint.questai.factory.ModelFactory;
import com.sprint.questai.model.dto.response.ChunkResponse;
import com.sprint.questai.model.entity.ChunkInfo;
import com.sprint.questai.model.entity.Index;
import com.sprint.questai.model.enums.NameEnums;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.spi.prompt.PromptTemplateFactory;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class KnowledgeService {
    @Autowired
    ElasticsearchClient client;
    @Autowired
    private ModelFactory factory;


    public String createKnowledge(String userId) throws IOException {
        Index index = new Index();
        index.setUserId(userId);
        index.setChunkIds(new ArrayList<>());
        index.setKbName("默认知识库");
        // random generate a kbId
        String kbId = RandomUtil.randomString(20);
        index.setKbid(kbId);
        IndexResponse index1 = client.index(s -> s.id(kbId).index(NameEnums.KB_INDEX.toString()).document(index));
        //
        return kbId;
    }
    public List<Index> getKnowledgesByUserId(String userId) throws IOException {
        SearchResponse<Index> response = client.search(s -> s.index(NameEnums.KB_INDEX.toString())
                .query(q -> q.match(m -> m.field("userId").query(userId))), Index.class);
        List<Hit<Index>> hits = response.hits().hits();
        return hits.stream().map(Hit::source).collect(ArrayList::new, ArrayList::add, ArrayList::addAll);
    }
    public String  addEmptyChunk(String knowledgeId) throws IOException {
        ChunkInfo emptyChunk = new ChunkInfo();
        Map<String,Object> metadata = new HashMap<>();
        metadata.put("knowledgeId",knowledgeId);
        emptyChunk.setMetadata(metadata);
        emptyChunk.setText("");
        String id=RandomUtil.randomString(20);
        client.index(s->
                s.id(id).index(NameEnums.generateEmbeddingIndexName(NameEnums.CHUNK_INDEX.toString(),NameEnums.DEFAULT_EMBEDDING_MODEL.toString()))
                .document(emptyChunk));
        refreshKBIndex(knowledgeId);
        return id;
    }

    public String addChunk(String knowledgeId, String chunk) throws IOException {
        //search the index by knowledgeId
        String s = updateChunkEmbedding(knowledgeId, chunk);
        refreshKBIndex(knowledgeId);
        return s;
    }
    public String addChunk(String knowledgeId, Document chunk) throws IOException {
        //search the index by knowledgeId
        String s = updateChunkEmbedding(knowledgeId, chunk);
        refreshKBIndex(knowledgeId);
        return s;
    }


    public String updateChunkByChunkId(String knowledgeId,String chunkId, String chunk) throws IOException {
        //search the index by knowledgeId
        if (chunkId != null) {
            deleteByChunkId(chunkId);
        }

        EmbeddingStoreIngestor ingestor =factory.createEmbeddingStoreIngestor(
                NameEnums.CHUNK_INDEX.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString(),
                NameEnums.DEFAULT_DOCUMENT_SPLITTER.toString()
        );
        Document document = Document.from(chunk);
        document.metadata().put("knowledgeId", knowledgeId);
        ingestor.ingest(document);
        String s=searchChunkIdByChunkText(knowledgeId, chunk);
        System.out.println(s);
        return chunkId;
    }
    public String updateChunkByChunkId(String knowledgeId,String chunkId, Document chunk) throws IOException {
        //search the index by knowledgeId
        if (chunkId != null) {
            deleteByChunkId(chunkId);
        }

        EmbeddingStoreIngestor ingestor =factory.createEmbeddingStoreIngestor(
                NameEnums.CHUNK_INDEX.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString(),
                NameEnums.ELASTIC_SEARCH_STORE.toString(),
                NameEnums.DEFAULT_DOCUMENT_SPLITTER.toString()
        );
        Document document =chunk;
        document.metadata().put("knowledgeId", knowledgeId);
        ingestor.ingest(document);
        String s=searchChunkIdByChunkText(knowledgeId, chunk);
        System.out.println(s);
        return chunkId;
    }
    /**
     * delete the chunk by knowledgeId
     * embed the new chunk
     * @param knowledgeId
     * @param chunk
     * @return
     * @throws IOException
     */
    public String updateChunkEmbedding(String knowledgeId, String chunk) throws IOException {
        //search the index by knowledgeId
        return updateChunkEmbedding(knowledgeId,Document.from(chunk));
    }
    public String updateChunkEmbedding(String knowledgeId, Document chunk) throws IOException {
        //search the index by knowledgeId
        String chunkId = searchChunkIdByChunkText(knowledgeId, chunk);
        return updateChunkByChunkId(knowledgeId,chunkId,chunk);
    }
    /**
     * search the chunkId by knowledgeId and chunktext
     * @param knowledgeId
     * @param chunk
     * @return
     * @throws IOException
     */
    public String  searchChunkIdByChunkText(String knowledgeId,String chunk) throws IOException {
        //search the index by text and knowledgeId
        return searchChunkIdByChunkText(knowledgeId,Document.from(chunk));
    }
    public String  searchChunkIdByChunkText(String knowledgeId,Document chunk) throws IOException {
        //search the index by text and knowledgeId
        SearchResponse<ChunkInfo> response = client.search(s -> s.index(NameEnums.generateEmbeddingIndexName(
                        NameEnums.CHUNK_INDEX.toString(),
                        NameEnums.DEFAULT_EMBEDDING_MODEL.toString())).query(
//                        q -> q.match(m ->
//                                m.field("metadata.knowledgeId").query(knowledgeId)
//                                        .field("text").query(chunk)
//                        )
//                    q->q.bool(b->b.must(c->c.matchPhrase(
//                            m->m.field(("metadata.knowledgeId")).query(knowledgeId)
//                                    .field("text").query(chunk)))
//                    )
//                q->q.term(m->m.field("metadata.knowledgeId").value(knowledgeId).field("text").value(chunk))
                        q->q.matchPhrase(m->m.field("metadata.knowledgeId").query(knowledgeId)
                                .field("text.keyword").query(chunk.text())
                        )
                ),
                ChunkInfo.class);
        List<Hit<ChunkInfo>> hits = response.hits().hits();
        if(hits.size()>0){
            //delete the old chunk
            return hits.get(0).id();
        }
        return null;
    }
    public void deleteByChunkId(String chunkId) throws IOException {
        client.delete(s -> s.index(NameEnums.generateEmbeddingIndexName(
                NameEnums.CHUNK_INDEX.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString())).id(chunkId));
        System.out.println("delete chunk"+chunkId);
    }


    /**
     * refresh the index by knowledgeId
     * search the index by knowledgeId and update the chunkIds
     * @param knowledgeId
     * @throws IOException
     */
    public void refreshKBIndex(String knowledgeId) throws IOException {
        Index index = searchKBIndexIdByKnowledgeId(knowledgeId);
        // 构建更新
        Index finalIndex = index;
        finalIndex.getChunkIds().clear();
        searchChunkByKnowledgeId(knowledgeId).forEach(hit -> {
            finalIndex.addChunkId(hit.id());
        });
        client.index(s -> s.id(knowledgeId).
                index(NameEnums.KB_INDEX.toString())
                .document(finalIndex));
        System.out.println("knowledgeId found");
    }
    public void addChunkIdToKBIndex(String knowledgeId,String chunkId) throws IOException {
        Index index = searchKBIndexIdByKnowledgeId(knowledgeId);
        index.addChunkId(chunkId);
        client.index(s -> s.id(knowledgeId).
                index(NameEnums.KB_INDEX.toString())
                .document(index));
    }
    public void deleteChunkIdFromKBIndex(String knowledgeId,String chunkId) throws IOException {
        Index index = searchKBIndexIdByKnowledgeId(knowledgeId);
        index.removeChunkId(chunkId);
        client.index(s -> s.id(knowledgeId).
                index(NameEnums.KB_INDEX.toString())
                .document(index));
    }
    /**
     * search the index by knowledgeId
     * @param knowledgeId
     * @return
     * @throws IOException
     */
    public Index searchKBIndexIdByKnowledgeId(String knowledgeId) throws IOException {
        SearchResponse<Index> kbid = client.
                search(s -> s.index(NameEnums.KB_INDEX.toString())
                                .query(q->q.match(m->m.field(NameEnums.KBID_NAME.toString()).query(knowledgeId)))
                        , Index.class);
        System.out.println(kbid);
        if (kbid.hits().hits().size() != 1) {
            System.out.println("knowledgeId found error  " + kbid.hits().hits().size());
            throw new RuntimeException("knowledgeId not found or found more than one");
        }
        return kbid.hits().hits().get(0).source();
    }
    public List<Hit<ChunkInfo>>  searchChunkByKnowledgeId(String knowledgeId) throws IOException {
        String chunkIndexName=NameEnums.generateEmbeddingIndexName(
                NameEnums.CHUNK_INDEX.toString(),
                NameEnums.DEFAULT_EMBEDDING_MODEL.toString()
        );
        SearchResponse<ChunkInfo> chunks = client.
                search(s -> s.index(chunkIndexName)
                                .query(q -> q.match(m -> m.field("metadata.knowledgeId")
                                        .query(knowledgeId))).size(50)
                        , ChunkInfo.class);
        System.out.println(chunks);
        return chunks.hits().hits();
    }
    public List<ChunkResponse> searchChunksByKnowledgeIdResponse(String knowledgeId) throws IOException {
        return searchChunkByKnowledgeId(knowledgeId).stream().map(hit -> {
            ChunkResponse response = new ChunkResponse();
            response.setKnowledgeId(knowledgeId);
            response.setChunkId(hit.id());
            response.setText(hit.source().getText());
            return response;
        }).collect(ArrayList::new, ArrayList::add, ArrayList::addAll);
    }

}

